import asyncio
import os
import dotenv

from openai import AsyncOpenAI

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
)

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)


CUSTOM_MODEL_PROVIDER = CustomModelProvider()
agent = Agent(
    name="数学辅导员",
    instructions="""
    你是一个数学辅导员，你的任务是帮助用户解决数学问题。
    每一步都要解释你的推理过程并包含例子。
    如果用户的问题是数学问题，请使用数学公式和定理来解决。
    如果用户的问题是其他问题，请回答他臣妾做不到啊
    """
)

async def main():
    user_input = input("请输入您的数学问题: ")
    result = await Runner.run(
        starting_agent = agent,
        input =user_input,
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    )
    print(result.final_output)

if __name__ == "__main__":  
    asyncio.run(main())








